#include <conio.h>
#include "cvxStereoIntrinsicSolver.h"

using cvx::Rgb8UImage;
using cvx::Bw8UImage;
using cvx::Bw32FImage;
using cvx::Bw32SImage;
using cvx::Rgb32FImage;

CvxStereoIntrinsicSolver::CvxStereoIntrinsicSolver()
{
	m_leftETree = NULL;  //solve unknown area of local solver result
	m_leftMTree = NULL;  //refine known area
	m_rightETree = NULL;
	m_rightMTree = NULL;

	m_leftIntrinsicData = NULL;
	m_rightIntrinsicData = NULL;
}
CvxStereoIntrinsicSolver::~CvxStereoIntrinsicSolver()
{
	//release data and images

}

void CvxStereoIntrinsicSolver::setImage(const Mat &leftImage, const Mat &rightImage, float threshold)
{
	assert(!leftImage.empty());
	assert(!rightImage.empty());
	assert(leftImage.rows == rightImage.rows);
	assert(leftImage.cols == rightImage.cols);

	m_leftIntrinsicData  = new CvxIntrinsicSolver();
	assert(m_leftIntrinsicData);
	
	m_leftIntrinsicData->setImage(leftImage);	
	m_leftIntrinsicData->computeEdgeMap(threshold);
	m_leftIntrinsicData->computeMiniShading();

	m_rightIntrinsicData = new CvxIntrinsicSolver();
	assert(m_rightIntrinsicData);
	m_rightIntrinsicData->setImage(rightImage);
	m_rightIntrinsicData->computeEdgeMap(threshold);
	m_rightIntrinsicData->computeMiniShading();

}

bool CvxStereoIntrinsicSolver::SetETree(const int stddevThreshold, const int sizeThreshold)
{
	//check input parameters	
	int h = m_leftIntrinsicData->m_orgImg.rows;
	int w = m_leftIntrinsicData->m_orgImg.cols;	

	unsigned int level = CvxQuadTree::GetLevel(std::max(w, h));

	//set left E tree
	m_leftETree = new CvxAdaptiveQuadTree(level, m_leftIntrinsicData->m_orgImg, stddevThreshold);
	m_leftETree->create(NULL);
	m_leftETree->balance();
	m_leftETree->adaptiveInitLeafData(m_leftIntrinsicData->m_orgImg, m_leftIntrinsicData->m_edgeMap, NULL, sizeThreshold);

	//set right E tree
	m_rightETree = new CvxAdaptiveQuadTree(level, m_rightIntrinsicData->m_orgImg, stddevThreshold);
	m_rightETree->create(NULL);
	m_rightETree->balance();
	m_rightETree->adaptiveInitLeafData(m_rightIntrinsicData->m_orgImg, m_rightIntrinsicData->m_edgeMap, NULL, sizeThreshold);
	m_rightETree->updateUnknownIndex(m_leftETree->unknownNodeNumber(), NULL);

	//test
	if(0)
	{
		m_leftETree->draw(NULL);
		m_rightETree->draw(NULL);
		cv::imshow("left E tree", m_leftETree->m_orgimg);
		cv::imshow("right E tree", m_rightETree->m_orgimg);
	}
	return true;
}

bool CvxStereoIntrinsicSolver::SetETree(const Mat &leftInitialShadingMap, const Mat &rightInitialShadingMap, const int threshold)
{
	return true;
}

bool CvxStereoIntrinsicSolver::MatchLeftRightETree(float rgbThreshold, const int sizeThreshold, bool isShow)
{
	//check input
	assert(m_leftETree && m_rightETree);

	m_leftETree->matchNode(rgbThreshold, sizeThreshold, NULL, m_rightETree);
	m_rightETree->matchNode(rgbThreshold, sizeThreshold, NULL, m_leftETree);

	m_leftETree->crossMatchNode(NULL, m_rightETree);
	m_rightETree->crossMatchNode(NULL, m_leftETree);

	if(isShow)
	{
		Mat resultImage;
		m_leftETree->drawMatchNode(NULL, m_rightETree, resultImage);

		cv::imshow("combined image", resultImage);
	//	cv::imwrite("combined.png", resultImage);
	}
	

	return true;
}

bool CvxStereoIntrinsicSolver::SolveETree(bool useEdge, bool useMatch, bool test)
{
	if (m_leftETree && m_rightETree)
	{
		vector<map<int, double>> leftMap;
		vector<double> rightVec;

		//equation from left image
		m_leftETree->planeEquation(NULL, leftMap, rightVec);
		if(useEdge)
		{
			m_leftETree->edgeEquation(NULL, leftMap, rightVec);
		}
		assert(leftMap.size() == rightVec.size());

		//equation from right image
		m_rightETree->planeEquation(NULL, leftMap, rightVec);
		if(useEdge)
		{
			m_rightETree->edgeEquation(NULL, leftMap, rightVec);
		}
		assert(leftMap.size() == rightVec.size());

		//equation from left <-> right match
		if(useMatch)
		{
			m_leftETree->matchNodeEquation(NULL, leftMap, rightVec);	
		}			

		int unknownNum = m_leftETree->unknownNodeNumber() + m_rightETree->unknownNodeNumber();
		_cprintf("unknow num is %d.\n", unknownNum);
		assert(unknownNum > 0);
		vector<double> result(unknownNum);
		bool isOk = LeastSquare::solver(leftMap, rightVec, true, unknownNum, &result[0]);
		if (isOk)
		{
			m_leftETree->setShadingResult(m_leftIntrinsicData->m_shading, NULL, result);
			m_rightETree->setShadingResult(m_rightIntrinsicData->m_shading, NULL, result);

			if(test)
			{
				Mat leftShowShading = m_leftIntrinsicData->m_shading.clone();				
				leftShowShading.convertTo(leftShowShading, CV_8U, 255.0);
				Mat rightShowShading = m_rightIntrinsicData->m_shading.clone();
				rightShowShading.convertTo(rightShowShading, CV_8U, 255.0);
				cv::imshow("left shading E tree", leftShowShading);
				cv::imshow("right shading E tree", rightShowShading);
			}
			return true;
		}
		else
		{
			_cprintf("error: CvxStereoIntrinsicSolver::SolveETree solver failed.\n");
			return false;
		}		
	}
	else
	{		
		_cprintf("error: CvxStereoIntrinsicSolver set E tree first.\n");
		return false;
	}
	return true;
}
bool CvxStereoIntrinsicSolver::CalcMTree(bool reverse)
{	
	assert(m_leftETree);
	assert(m_rightETree);
	
	//calculate left M tree
	if (m_leftMTree == NULL)
	{
		m_leftMTree = new CvxAdaptiveQuadTree(*m_leftETree);
		if (reverse)
		{
			m_leftMTree->reverse(NULL);
		}
	}

	//calculate right M tree
	if(m_rightMTree == NULL)
	{
		m_rightMTree = new CvxAdaptiveQuadTree(*m_rightETree);
		if(reverse)
		{
			m_rightMTree->reverse(NULL);
		}
	}

	//update index
	m_rightMTree->updateUnknownIndex(m_leftMTree->unknownNodeNumber(), NULL);
	
	return true;
}
bool CvxStereoIntrinsicSolver::SolveMtree(bool useEdge, bool test)
{
	if (m_leftMTree && m_rightMTree)
	{
		vector<map<int, double>> leftMap;
		vector<double> rightVec;

		// equation from left image
		m_leftMTree->planeEquation(NULL, leftMap, rightVec);
		if(useEdge)
		{
			m_leftMTree->edgeEquation(NULL, leftMap, rightVec);
		}
		assert(leftMap.size() == rightVec.size());

		// equation from right image
		m_rightMTree->planeEquation(NULL, leftMap, rightVec);
		if(useEdge)
		{
			m_rightMTree->edgeEquation(NULL, leftMap, rightVec);
		}
		assert(leftMap.size() == rightVec.size());

		int unknownNum = m_leftMTree->unknownNodeNumber() + m_rightMTree->unknownNodeNumber();
		_cprintf("M tree unknownNum is %d.\n", unknownNum);

		vector<double> result(unknownNum);

		bool isOk = LeastSquare::solver(leftMap, rightVec, true, unknownNum, &result[0]);
		if (isOk)
		{
			m_leftMTree->setShadingResult(m_leftIntrinsicData->m_shading, NULL, result);
			m_rightMTree->setShadingResult(m_rightIntrinsicData->m_shading, NULL, result);

			if(test)
			{
				Mat leftShowShading = m_leftIntrinsicData->m_shading.clone();				
				leftShowShading.convertTo(leftShowShading, CV_8U, 255.0);
				Mat rightShowShading = m_rightIntrinsicData->m_shading.clone();
				rightShowShading.convertTo(rightShowShading, CV_8U, 255.0);
				cv::imshow("left shading M tree", leftShowShading);
				cv::imshow("right shading M tree", rightShowShading);
			}
			return true;
		}
		else
		{
			_cprintf("error: CvxStereoIntrinsicSolver::SolveMTree solver failed.\n");
			return false;
		}		
	}
	return false;	
}
	//up sample using A vector to get final result
bool CvxStereoIntrinsicSolver::UpsampleAvector(bool test)
{
	//check input
	assert(m_leftMTree);
	assert(m_rightMTree);

	//upsample A vector in left image
	bool isOk = upSampleAvectorInOneTree(m_leftMTree, m_leftIntrinsicData, false);
	if(!isOk)
	{
		_cprintf("error, left image upsamle A vector failed!\n");
		return false;
	}	

	//upsample A vector in right image
	m_rightMTree->updateUnknownIndex(-m_leftMTree->unknownNodeNumber(), NULL);
	isOk = upSampleAvectorInOneTree(m_rightMTree, m_rightIntrinsicData, false);
	if(!isOk)
	{
		_cprintf("error, right image upsamle A vector failed!\n");
		return false;
	}

	if(test)
	{
		Mat leftShowShading = m_leftIntrinsicData->m_shading.clone();				
		leftShowShading.convertTo(leftShowShading, CV_8U, 255.0);
		Mat rightShowShading = m_rightIntrinsicData->m_shading.clone();
		rightShowShading.convertTo(rightShowShading, CV_8U, 255.0);
		cv::imshow("left shading After up sampling A vector", leftShowShading);
		cv::imshow("right shading After up sampling A vector", rightShowShading);

		

		cv::imshow("left after up sampling A vector", m_leftIntrinsicData->m_albedo);
		cv::imshow("right after up sampling A vector", m_rightIntrinsicData->m_albedo);

		cv::imwrite("left_shading.png", leftShowShading); 
		cv::imwrite("right_shading.png", rightShowShading);

		cv::imwrite("left_reflectance.png", m_leftIntrinsicData->m_albedo);
		cv::imwrite("right_reflectance.png", m_rightIntrinsicData->m_albedo);
	}

	return true;

}

bool CvxStereoIntrinsicSolver::upSampleAvectorInOneTree(CvxAdaptiveQuadTree * const mTree, CvxIntrinsicSolver * const intrinsicData, bool test)
{
	assert(mTree);
	assert(intrinsicData);

	int h = intrinsicData->m_orgImg.rows;
	int w = intrinsicData->m_orgImg.cols;	
	int unknownNum = mTree->m_nodeNum;
	
	//calculate A vector
	vector<map<int, double>> leftMap;
	vector<double> rightVec;
	mTree->AvectorEquation(NULL, leftMap, rightVec);	

	vector<double> result(unknownNum*3);
	bool isOk = LeastSquare::solver(leftMap, rightVec, true, unknownNum*3, &result[0]);
	if (isOk)
	{
		Mat AvectorMatrix = Mat(h, w, CV_32FC3, Scalar(0, 0, -255, 0));
		mTree->setAvector(NULL, result);
		mTree->InterpolateAvector(NULL, AvectorMatrix);
		
		Bw32FImage sha_map(&intrinsicData->m_shading);
		Rgb32FImage avector_map(&AvectorMatrix);
		Rgb8UImage org_img(&intrinsicData->m_orgImg);
		Rgb8UImage albedo_img(&intrinsicData->m_albedo);
		Bw32FImage mini_map(&intrinsicData->m_miniShading);

		for (int y = 0; y<h; ++y)
		{
			for (int x =0; x<w; ++x)
			{
				sha_map[y][x] = org_img[y][x].r * avector_map[y][x].r + 
						org_img[y][x].g * avector_map[y][x].g + 
						org_img[y][x].b * avector_map[y][x].b;
				sha_map[y][x] = std::min(0.99f, sha_map[y][x]);
				if (sha_map[y][x] < mini_map[y][x])
				{
					sha_map[y][x] = mini_map[y][x];
				}
				albedo_img[y][x].r = org_img[y][x].r/sha_map[y][x];
				albedo_img[y][x].g = org_img[y][x].g/sha_map[y][x];
				albedo_img[y][x].b = org_img[y][x].b/sha_map[y][x];			
			}
		}

		if(test)
		{
			

		}
		return true;
	}
	return false;

}


